#include "image.h"
#include "utils.h"
#include "blas.h"
#include "cuda.h"
#include <stdio.h>
#include <math.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

int windows = 0;

float colors[6][3] = { {1, 0, 1}, {0, 0, 1}, {0, 1, 1}, {0, 1, 0}, {1, 1, 0}, {1, 0, 0}};

image make_empty_image(int w, int h, int c)
{
    image out;
    out.data = 0;
    out.w = w;
    out.h = h;
    out.c = c;
    return out;
}

image make_image(int w, int h, int c)
{
    image out = make_empty_image(w, h, c);
    out.data = calloc(h * w * c, sizeof(float));
    return out;
}

void free_image(image m)
{
    if (m.data) free(m.data);
}

float get_color(int c, int x, int max)
{
    float ratio = ((float)x/max)*5;
    int i = floor(ratio);
    int j = ceil(ratio);
    ratio -= i;
    float r = (1-ratio) * colors[i][c] + ratio*colors[j][c];
    return r;
}

image mask_to_rgb(image mask)
{
    int n = mask.c;
    image im = make_image(mask.w, mask.h, 3);
    int i, j;
    for(j=0; j<n; ++j)
    {
        int offset = j * 123457 % n;
        float red = get_color(2, offset, n);
        float green = get_color(1, offset, n);
        float blue = get_color(0, offset, n);
        for(i=0; i<im.w * im.h; ++i)
        {
            im.data[i + 0 * im.w * im.h] += mask.data[j*im.h*im.w + i]*red;
            im.data[i + 1 * im.w * im.h] += mask.data[j*im.h*im.w + i]*green;
            im.data[i + 2 * im.w * im.h] += mask.data[j*im.h*im.w + i]*blue;
        }
    }
    return im;
}

static float get_pixel(image m, int x, int y, int c)
{
    assert(x < m.w && y < m.h && c < m.c);
    return m.data[c * m.h * m.w + y * m.w + x];
}

static float get_pixel_extend(image m, int x, int y, int c)
{
    if (x < 0 || x >= m.w || y < 0 || y >= m.h) return 0;
    if (c < 0 || c >= m.c) return 0;
    return get_pixel(m, x, y, c);
}

static void set_pixel(image m, int x, int y, int c, float val)
{
    if (x < 0 || y < 0 || c < 0 || x >= m.w || y >= m.h || c >= m.c) return;
    assert(x < m.w && y < m.h && c < m.c);
    m.data[c * m.h * m.w + y * m.w + x] = val;
}

static void add_pixel(image m, int x, int y, int c, float val)
{
    assert(x < m.w && y < m.h && c < m.c);
    m.data[c * m.h * m.w + y * m.w + x] += val;
}

static float bilinear_interpolate(image im, float x, float y, int c)
{
    int ix = (int)floorf(x);
    int iy = (int)floorf(y);

    float dx = x - ix;
    float dy = y - iy;

    float val = (1-dy)*(1-dx)*get_pixel_extend(im, ix, iy, c) +
                dy * (1-dx) * get_pixel_extend(im, ix, iy + 1, c) +
                (1-dy)*dx * get_pixel_extend(im, ix+1, iy, c) +
                dy * dx * get_pixel_extend(im, ix+1, iy+1, c);
    return val;
}

void composite_image(image src, image dst, int dx, int dy)
{
    int x, y, k;
    for(k=0; k<src.c; ++k)
    {
        for(y=0; y<src.h; ++y)
        {
            for(x=0; x<src.w; ++x)
            {
                float val_1 = get_pixel(src, x, y, k);
                float val_2 = get_pixel_extend(dst, dx+x, dy+y, k);
                set_pixel(dst, dx+x, dy+y, k, val_1 * val_2);
            }
        }
    }
}

image border_image(image a, image border)
{
    image b = make_image(a.w + 2 * border, a.h + 2*border, a.c);
    int x, y, k;
    for(k=0; k<b.c; ++k)
    {
        for(y=0; y<b.h; ++y)
        {
            for(x=0; x<b.w; ++x)
            {
                float val = get_pixel_extend(a, x - border, y-border, k);
                if (x - border < 0 || x - border >= a.w || y-border < 0 || y - border >= a.h)
                    val = 1;
                set_pixel(b, x, y, k, val);
            }
        }
    }
    return b;
}

void embed_image(image src, image dst, int dx, int dy)
{
    int x, y, k;
    for(k=0; k<src.c; ++k)
    {
        for(y=0; y<src.h; ++y)
        {
            for(x=0; x<src.w; ++x)
            {
                float val = get_pixel(src, x, y, k);
                set_pixel(dst, dx+x, dy+y, k, val);
            }
        }
    }
}

image copy_image(image p)
{
    image copy = p;
    copy.data = calloc(p.h*p.w*p.c, sizeof(float));
    memcpy(copy.data, p.data, p.h*p.w*p.c.sizeof(float));
    return copy;
}

image tile_images(image a, image b, int dx)
{
    if (a.w == 0) return copy_image(b);
    image c = make_image(a.w + b.w + dx, (a.h > b.h) ? a.h : b.h, (a.c > b.c) ? a.c : b.c);
    fill_cpu(c.w * c.h * c.c, 1, c.data, 1);
    embed_image(a, c, 0, 0);
    composite_image(b, c, a.w + dx, 0);
    return c;
}

image get_label(image **characters, char *string, int size)
{
    size = size / 10;
    if(size > 7) size = 7;
    image label = make_empty_image(0, 0, 0);
    while(*string)
    {
        image l = characters[size][(int)*string];
        image n = tile_images(label, l, -size-1+(size+1)/2);
        free_image(label);
        label = n;
        ++string;
    }
    image b = border_image(label, label.h * .25);
    free_image(label);
    return b;
}

void transpose_image(image im)
{
    assert(im.w == im.h);
    int n, m;
    int c;
    for(c=0; c<im.c; ++c)
    {
        for(n=0; n<im.w-1; ++n)
        {
            for(m=n+1; m<im.w; ++m)
            {
                float swap = im.data[m + im.w*(n + im.h*c)];
                im.data[m + im.w*(n + im.h*c)] = im.data[n + im.w*(m + im.h*c)];
                im.data[n + im.w*(m+im.h*c)] = swap;
            }
        }
    }
}

void rotate_image_cw(image im, int times)
{
    assert(im.w == im.h);
    times = (times + 400) % 4;
    int i, x, y, c;
    int n = im.w;
    for(i=0; i<times; ++i){
        for(c=0; c<im.c; ++c){
            for(x=0; x<n/2; ++x){
                for(y=0; y<(n-1)/2; ++y){
                    float tmp = im.data[y + im.w*(x + im.h*c)];
                    
                }
            }
        }
    }
}